#ifndef QUADTREE_H
#define QUADTREE_H

#include <cstdlib>

#include <QLineF>
#include <QRectF>
#include <QSet>
#include <QSizeF>
#include <QStack>
#include <QVector3D>

#include "geometry/edge3d.h"
#include "geometry/sphere.h"
#include "object/face.h"

template< class C >
class Quadtree : protected QRectF
{
public:
    Quadtree(const QRectF& rkRegion = QRectF(-200, -200, 400, 400), quint8 uiMaxDepth = 5, quint8 uiMaxCount = 4);
    Quadtree(qreal x, qreal y, qreal width, qreal height, quint8 uiMaxDepth = 5, quint8 uiMaxCount = 4);
    ~Quadtree();

    Quadtree* root() const;
    const Quadtree* parent() const;
    QSet< const Quadtree* > children() const;
    QSet< const Quadtree* > leaves() const;

    quint8 depth() const;
    quint32 count() const;

    QRectF bounds() const;
    void setBounds(const QRectF& rkBounds);
    void setBounds(qreal x, qreal y, qreal width, qreal height);

    const QSet< const C* >& items() const;

    void add(const C& rkItem);
    bool remove(const C& rkItem);
    void clear();

    const Quadtree* findItem(const C& rkItem) const;
    const Quadtree* findAt(const QVector3D& rkPosition) const;

protected:
    Quadtree(qreal x, qreal y, qreal width, qreal height, quint8 uiMaxDepth, quint8 uiMaxCount, Quadtree& rkParent);
    Quadtree(const Quadtree&);
    Quadtree& operator = (const Quadtree&);

private:
    typedef enum
    {
        Inside,
        Overlapping,
        Outside
    } Layout;

    Layout layout(const C& rkItem) const;

    void accept(const C& rkItem);
    void collapse();
    void expand();
    void spread(const C& rkItem);

    //------------------------------------------------
    //  Attributes.
    //------------------------------------------------
    // Parameters.
    static const quint8 s_uiPadding = 1;

    const quint8        m_uiMaxDepth;
    const quint8        m_uiMaxCount;

    // Hierarchy.
    Quadtree *const     m_pkParent;
    Quadtree*           m_akChildren;

    // Container.
    QSet< const C* >    m_kItems;
};

template< class C >
Quadtree< C >::Quadtree(const QRectF &rkRegion, quint8 uiMaxDepth, quint8 uiMaxCount)
    : QRectF(rkRegion),
      m_uiMaxDepth(uiMaxDepth),
      m_uiMaxCount(uiMaxCount),
      m_pkParent(NULL),
      m_akChildren(NULL)
{
}

template< class C >
Quadtree< C >::Quadtree(qreal x, qreal y, qreal width, qreal height, quint8 uiMaxDepth, quint8 uiMaxCount)
    : QRectF(x, y, width, height),
      m_uiMaxDepth(uiMaxDepth),
      m_uiMaxCount(uiMaxCount),
      m_pkParent(NULL),
      m_akChildren(NULL)
{
}

template< class C >
Quadtree< C >::~Quadtree()
{
    // Nothing to do if the node is not expanded.
    if(!m_akChildren)
        return;

    // Delete children.
    collapse();
}

template< class C >
Quadtree< C >* Quadtree< C >::root() const
{
    const Quadtree* pkCurrent = this;

    while(pkCurrent->m_pkParent)
        pkCurrent = pkCurrent->m_pkParent;

    return const_cast< Quadtree* >(pkCurrent);
}

template< class C >
const Quadtree< C >* Quadtree< C >::parent() const
{
    return m_pkParent;
}

template< class C >
QSet< const Quadtree< C >* > Quadtree< C >::children() const
{
    QSet< const Quadtree* > result;

    if(m_akChildren)
    {
        for(quint8 i = 0; i < 4; ++i)
        {
            result.insert(m_akChildren + i);
        }
    }

    return result;
}

template< class C >
QSet< const Quadtree< C >* > Quadtree< C >::leaves() const
{
    QSet< const Quadtree* > kLeaves;
    QStack< const Quadtree* > kStack;
    kStack.push(this);

    while(!kStack.empty())
    {
        const Quadtree* pkCurrent = kStack.pop();

        // Push the next children.
        if(pkCurrent->m_akChildren)
        {
            for(quint8 i = 0; i < 4; ++i)
                kStack.push(pkCurrent->m_akChildren + i);
        }
        // If no children then it's a leaf.
        else
        {
            kLeaves.insert(pkCurrent);
        }
    }

    return kLeaves;
}

template< class C >
quint8 Quadtree< C >::depth() const
{
    const Quadtree* pkCurrent = this;
    quint8 uiDepth = 0;

    while(pkCurrent->m_pkParent)
    {
        ++uiDepth;
        pkCurrent = pkCurrent->m_pkParent;
    }

    return uiDepth;
}

template< class C >
quint32 Quadtree< C >::count() const
{
    return m_kItems.count();
}

template< class C >
QRectF Quadtree< C >::bounds() const
{
    return *this;
}

template< class C >
void Quadtree< C >::setBounds(const QRectF& rkBounds)
{
    // Rebuild.
    setBounds(rkBounds.x(), rkBounds.y(), rkBounds.width(), rkBounds.height());
}

template< class C >
void Quadtree< C >::setBounds(qreal x, qreal y, qreal width, qreal height)
{
    // Only root can be resized.
    Q_ASSERT(!m_pkParent);

    // Collapse if needed.
    if(m_akChildren)
        collapse();

    // Actually setup the new bounds and
    // ensure every item is accepted.
    setRect(x, y, width, height);

    foreach(const C* pkItem, m_kItems)
    {
        if(layout(*pkItem) != Inside)
            accept(*pkItem);
    }

    // Expand the quadtree if needed and spread the items.
    if(!m_akChildren && !(count() < m_uiMaxCount))
    {
        expand();

        // Check that the quadtree was able to expand.
        if(m_akChildren)
            foreach(const C* pkItem, m_kItems)
                spread(*pkItem);
    }
}

template< class C >
const QSet< const C * >& Quadtree< C >::items() const
{
    return m_kItems;
}

template< class C >
void Quadtree< C >::add(const C& rkItem)
{
    // If this is root, adjust boundaries to accept the new item.
    if(!m_pkParent && layout(rkItem) != Inside)
        accept(rkItem);

    Q_ASSERT(layout(rkItem) != Outside);
    m_kItems.insert(&rkItem);

    // Expand the quadtree if needed and spread the items.
    if(!m_akChildren && !(count() < m_uiMaxCount))
    {
        expand();

        // Check that the quadtree was able to expand.
        if(m_akChildren)
            foreach(const C* pkItem, m_kItems)
                spread(*pkItem);
    }

    // Spread the item if the quadtree has children.
    if(m_akChildren)
        spread(rkItem);
}

template< class C >
bool Quadtree< C >::remove(const C& rkItem)
{
    // The item doesn't belong to the quadtree.
    if(m_kItems.find(&rkItem) == m_kItems.end())
        return false;

    // Proceed to DFS removal.
    QStack< Quadtree< C >* > kStack;
    kStack.push(this);

    while(!kStack.empty())
    {
        Quadtree< C >* pkCurrent = kStack.pop();
        typename QSet< const C* >::Iterator kIt = pkCurrent->m_kItems.find(&rkItem);

        // Stop removal on this branch.
        if(kIt == pkCurrent->m_kItems.end())
            continue;

        // Push the next children.
        if(pkCurrent->m_akChildren)
        {
            // Current quadtree is collapsing.
            if(pkCurrent->count() - 1 < m_uiMaxCount)
            {
                pkCurrent->collapse();
            }
            // Otherwise push the children.
            else
            {
                for(quint8 i = 0; i < 4; ++i)
                    kStack.push(pkCurrent->m_akChildren + i);
            }
        }

        // Finally erase from set.
        pkCurrent->m_kItems.erase(kIt);
    }

    return true;
}

template< class C >
void Quadtree< C >::clear()
{
    // Collapse if needed.
    if(m_akChildren)
        collapse();

    m_kItems.clear();
}

template< class C >
const Quadtree< C >* Quadtree< C >::findItem(const C& rkItem) const
{
    const Quadtree* pkCurrent = this;

    while(pkCurrent->m_akChildren)
    {
        quint8 i = 0;
        for(i = 0; i < 4; ++i)
        {
            Layout kLayout = pkCurrent->m_akChildren[i].layout(rkItem);

            // The item is inside a child quadtree.
            // Proceed with the child quadtree.
            if(kLayout == Inside)
            {
                pkCurrent = pkCurrent->m_akChildren + i;
                break;
            }

            // The current quadtree is the smallest containing the item.
            if(kLayout == Overlapping)
                return pkCurrent;
        }

        // Failed to find corresponding quadtree.
        if(i == 4)
            return NULL;
    }

    // Found the deepest quadtree containing the item.
    return pkCurrent;
}

template< class C >
const Quadtree< C >* Quadtree< C >::findAt(const QVector3D& rkPosition) const
{
    const Quadtree* pkCurrent = this;

    // Check if base quadtree contains the projected vertex.
    if(!pkCurrent->contains(rkPosition.x(), rkPosition.z()))
        return NULL;

    // Proceed on children.
    while(pkCurrent->m_akChildren)
    {
        quint8 i = 0;
        for(i = 0; i < 4; ++i)
        {
            if(pkCurrent->m_akChildren[i].contains(rkPosition.x(), rkPosition.z()))
            {
                pkCurrent = pkCurrent->m_akChildren + i;
                break;
            }
        }

        // Found the deepest quadtree containing the projected vertex.
        if(i == 4)
            break;
    }

    return pkCurrent;
}

template< class C >
Quadtree< C >::Quadtree(qreal x, qreal y, qreal width, qreal height,
                        quint8 uiMaxDepth, quint8 uiMaxCount,
                        Quadtree& rkParent)
    : QRectF(x, y, width, height),
      m_uiMaxDepth(uiMaxDepth),
      m_uiMaxCount(uiMaxCount),
      m_pkParent(&rkParent),
      m_akChildren(NULL)
{
}

template< class C >
typename Quadtree< C >::Layout Quadtree< C >::layout(const C& rkItem) const
{
    Q_ASSERT(rkItem.isValid());
    Q_ASSERT(isValid());

    // Inside the quadtree.
    if(contains(rkItem))
        return Inside;

    // Overlapping multiple quadtrees.
    if(intersected(rkItem).isValid())
        return Overlapping;

    // Outside the quadtree.
    return Outside;
}

// Sphere specialization.
template<>
inline Quadtree< Sphere >::Layout Quadtree< Sphere >::layout(const Sphere& rkSphere) const
{
    Q_ASSERT(rkSphere.isValid());
    Q_ASSERT(isValid());

    qreal radius = rkSphere.radius();
    QRectF kBounds(rkSphere.center().x() - radius,
                   rkSphere.center().z() - radius,
                   radius * 2,
                   radius * 2);

    kBounds.adjust(-s_uiPadding, -s_uiPadding, s_uiPadding, s_uiPadding);

    // Inside the quadtree.
    if(contains(kBounds))
        return Inside;

    // Overlapping multiple quadtrees.
    if(intersected(kBounds).isValid())
        return Overlapping;

    // Outside the quadtree.
    return Outside;
}

// Edge3D specialization.
template<>
inline Quadtree< Edge3D >::Layout Quadtree< Edge3D >::layout(const Edge3D& rkEdge) const
{
    Q_ASSERT(rkEdge.isValid());
    Q_ASSERT(isValid());

    QRectF kBounds(qMin(rkEdge.p1().x(), rkEdge.p2().x()),
                   qMin(rkEdge.p1().z(), rkEdge.p2().z()),
                   qAbs(rkEdge.p1().x() - rkEdge.p2().x()),
                   qAbs(rkEdge.p1().z() - rkEdge.p2().z()));

    kBounds.adjust(-s_uiPadding, -s_uiPadding, s_uiPadding, s_uiPadding);

    // Inside the quadtree.
    if(contains(kBounds))
        return Inside;

    // Overlapping multiple quadtrees.
    if(intersected(kBounds).isValid())
        return Overlapping;

    // Outside the quadtree.
    return Outside;
}

// Face specialization.
template<>
inline Quadtree< Face >::Layout Quadtree< Face >::layout(const Face& rkFace) const
{
    Q_ASSERT(rkFace.isValid());
    Q_ASSERT(isValid());

    qreal xMin = 0, xMax = 0;
    qreal zMin = 0, zMax = 0;

    foreach(const QVector3D* pkVertex, rkFace.vertices())
    {
        xMin = qMin(xMin, static_cast< qreal >(pkVertex->x()));
        xMax = qMax(xMax, static_cast< qreal >(pkVertex->x()));

        zMin = qMin(zMin, static_cast< qreal >(pkVertex->z()));
        zMax = qMax(zMax, static_cast< qreal >(pkVertex->z()));
    }

    QRectF kBounds = QRect(xMin, zMin, xMax - xMin, zMax - zMin);

    kBounds.adjust(-s_uiPadding, -s_uiPadding, s_uiPadding, s_uiPadding);

    // Inside the quadtree.
    if(contains(kBounds))
        return Inside;

    // Overlapping multiple quadtrees.
    if(intersected(kBounds).isValid())
        return Overlapping;

    // Outside the quadtree.
    return Outside;
}

// QLineF specialization.
template<>
inline Quadtree< QLineF >::Layout Quadtree< QLineF >::layout(const QLineF& rkLine) const
{
    Q_ASSERT(isValid());

    QRectF kBounds(qMin(rkLine.p1().x(), rkLine.p2().x()),
                   qMin(rkLine.p1().y(), rkLine.p2().y()),
                   rkLine.dx(),
                   rkLine.dy());

    kBounds.adjust(-s_uiPadding, -s_uiPadding, s_uiPadding, s_uiPadding);

    // Inside the quadtree.
    if(contains(kBounds))
        return Inside;

    // Overlapping multiple quadtrees.
    if(intersected(kBounds).isValid())
        return Overlapping;

    // Outside the quadtree.
    return Outside;
}

template< class C >
void Quadtree< C >::accept(const C& rkItem)
{
    Q_ASSERT(!m_pkParent && layout(rkItem) != Inside);

    QRectF kBounds = *this | rkItem;

    kBounds.adjust(-s_uiPadding, -s_uiPadding, s_uiPadding, s_uiPadding);

    setRect(kBounds.x(), kBounds.y(),
            kBounds.width(), kBounds.height());

    // Collapse if it has children.
    if(m_akChildren)
        collapse();
}

// Sphere specialization.
template<>
inline void Quadtree< Sphere >::accept(const Sphere& rkSphere)
{
    Q_ASSERT(!m_pkParent && layout(rkSphere) != Inside);

    qreal radius = rkSphere.radius();
    QRectF kBounds = *this | QRectF(rkSphere.center().x() - radius,
                                    rkSphere.center().z() - radius,
                                    radius * 2, radius * 2);

    kBounds.adjust(-s_uiPadding, -s_uiPadding, s_uiPadding, s_uiPadding);

    setRect(kBounds.x(), kBounds.y(),
            kBounds.width(), kBounds.height());

    // Collapse if it has children.
    if(m_akChildren)
        collapse();
}

// Edge3D specialization.
template<>
inline void Quadtree< Edge3D >::accept(const Edge3D& rkEdge)
{
    Q_ASSERT(!m_pkParent && layout(rkEdge) != Inside);

    QRectF kBounds = *this | QRectF(qMin(rkEdge.p1().x(), rkEdge.p2().x()),
                                    qMin(rkEdge.p1().z(), rkEdge.p2().z()),
                                    qAbs(rkEdge.p1().x() - rkEdge.p2().x()),
                                    qAbs(rkEdge.p1().z() - rkEdge.p2().z()));

    kBounds.adjust(-s_uiPadding, -s_uiPadding, s_uiPadding, s_uiPadding);

    setRect(kBounds.x(), kBounds.y(),
            kBounds.width(), kBounds.height());

    // Collapse if it has children.
    if(m_akChildren)
        collapse();
}

// Face specialization.
template<>
inline void Quadtree< Face >::accept(const Face& rkFace)
{
    Q_ASSERT(!m_pkParent && layout(rkFace) != Inside);

    qreal xMin = 0, xMax = 0;
    qreal zMin = 0, zMax = 0;

    foreach(const QVector3D* pkVertex, rkFace.vertices())
    {
        xMin = qMin(xMin, static_cast< qreal >(pkVertex->x()));
        xMax = qMax(xMax, static_cast< qreal >(pkVertex->x()));

        zMin = qMin(zMin, static_cast< qreal >(pkVertex->z()));
        zMax = qMax(zMax, static_cast< qreal >(pkVertex->z()));
    }

    QRectF kBounds = *this | QRect(xMin, zMin, xMax - xMin, zMax - zMin);

    kBounds.adjust(-s_uiPadding, -s_uiPadding, s_uiPadding, s_uiPadding);

    setRect(kBounds.x(), kBounds.y(),
            kBounds.width(), kBounds.height());

    // Collapse if it has children.
    if(m_akChildren)
        collapse();
}

// QLineF specialization.
template<>
inline void Quadtree< QLineF >::accept(const QLineF& rkLine)
{
    Q_ASSERT(!m_pkParent && layout(rkLine) != Inside);

    QRectF kBounds = *this | QRectF(qMin(rkLine.p1().x(), rkLine.p2().x()),
                                    qMin(rkLine.p1().y(), rkLine.p2().y()),
                                    rkLine.dx(),
                                    rkLine.dy());

    kBounds.adjust(-s_uiPadding, -s_uiPadding, s_uiPadding, s_uiPadding);

    setRect(kBounds.x(), kBounds.y(),
            kBounds.width(), kBounds.height());

    // Collapse if it has children.
    if(m_akChildren)
        collapse();
}

template< class C >
void Quadtree< C >::collapse()
{
    // Only meaningful to collapse an extanded quadtree.
    Q_ASSERT(m_akChildren);

    // Delete children.
    for(quint8 i = 0; i < 4; ++i)
        m_akChildren[i].~Quadtree();

    free(m_akChildren);
    m_akChildren = NULL;
}

template< class C >
void Quadtree< C >::expand()
{
    // Only leaves can be expanded.
    Q_ASSERT(!m_akChildren);

    // Do nothing if quadtree's depth is equal or over the limit.
    if(depth() >= m_uiMaxDepth)
        return;

    // Create children and initialize.
    m_akChildren = static_cast< Quadtree* >(malloc(sizeof(Quadtree) * 4));
    Q_CHECK_PTR(m_akChildren);

    QSizeF kChildSize = size() / 2;

    for(quint8 i = 0; i < 4; ++i)
    {
        new (m_akChildren + i) Quadtree< C >(x() + kChildSize.width() * (i % 2),
                                             y() + kChildSize.height() * (i / 2),
                                             kChildSize.width(),
                                             kChildSize.height(),
                                             m_uiMaxDepth,
                                             m_uiMaxCount,
                                             *this);
    }
}

template< class C >
void Quadtree< C >::spread(const C& rkItem)
{
    Q_ASSERT(layout(rkItem) != Outside);
    Q_ASSERT(m_akChildren);

    for(quint8 i = 0; i < 4; ++i)
        if(m_akChildren[i].layout(rkItem) != Outside)
            m_akChildren[i].add(rkItem);
}


#endif // QUADTREE_H
